package com.oliver.proxy.midjourney.application.service;

import com.oliver.proxy.midjourney.domain.entity.Prod;
import com.oliver.proxy.midjourney.enums.TaskStatus;
import com.oliver.proxy.midjourney.infrastructure.mapper.ProdMapper;
import com.oliver.proxy.midjourney.interfaces.controller.SubmitController;
import com.oliver.proxy.midjourney.interfaces.controller.TaskController;
import com.oliver.proxy.midjourney.interfaces.dto.SubmitImagineDTO;
import com.oliver.proxy.midjourney.interfaces.dto.TaskConditionDTO;
import com.oliver.proxy.midjourney.result.SubmitResultVO;
import com.oliver.proxy.midjourney.support.Task;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.PostConstruct;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.net.URL;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

@Slf4j
@Service
public class ImageProcessingService {
    private static final String BASE_PATH = "/Users/oliver/pic/preGenerate/";
    private static final String SEL_PATH = "/Users/oliver/pic/selectByDev/";
    private static final long IMAGINE_REQUEST_DELAY_MS = 5000; // 5 seconds delay between imagine requests
    private final ProdMapper prodMapper;
    private final TaskController taskController;
    private final SubmitController submitController;
    private final AtomicBoolean isProcessing = new AtomicBoolean(false);


    @Autowired
    public ImageProcessingService(ProdMapper prodMapper, TaskController taskController, SubmitController submitController) {
        this.prodMapper = prodMapper;
        this.taskController = taskController;
        this.submitController = submitController;
        // 确保目录存在
        createDirectoryIfNotExists();
    }

    private void createDirectoryIfNotExists() {
        File directory = new File(BASE_PATH);
        if (!directory.exists()) {
            directory.mkdirs();
        }
        File sel = new File(SEL_PATH);
        if (!sel.exists()) {
            sel.mkdirs();
        }
    }

    @PostConstruct
    public void init() {
        CompletableFuture.runAsync(this::processImages);
    }

    @Scheduled(fixedRate = 2000)
    public void checkAndRestartProcessing() {
        if (!isProcessing.get()) {
            log.info("Restarting Image processing...");
            CompletableFuture.runAsync(this::processImages);
        }
    }


    public void processImages() {
        try {
            if (!isProcessing.compareAndSet(false, true)) {
                return;
            }

            log.info("Starting Image processing in background thread");
            ImageIO.scanForPlugins();

            while (true) {

                // 获取所有需要处理的记录
                List<Prod> unprocessedProds = prodMapper.selectUnprocessedImages();
                if (unprocessedProds.isEmpty()) {
                    Thread.sleep(2000);
                    continue;
                }

                // 获取唯一的taskIds
                Set<String> uniqueTaskIds = unprocessedProds.stream()
                        .map(Prod::getAttritube1)
                        .filter(Objects::nonNull)
                        .collect(Collectors.toSet());

                // 创建TaskConditionDTO
                TaskConditionDTO conditionDTO = new TaskConditionDTO();
                conditionDTO.setIds(new ArrayList<>(uniqueTaskIds));

                // 调用taskController获取任务信息
                List<Task> tasks = taskController.listByIds(conditionDTO);
                Map<String, String> taskIdToImageUrl = new HashMap<>();

                List<Task> failureTask = new ArrayList<>();

                // 创建taskId到imageUrl的映射
                for (Task task : tasks) {
                    if (task.getStatus() == TaskStatus.SUCCESS && task.getImageUrl() != null && !task.getImageUrl().isEmpty()) {
                        taskIdToImageUrl.put(task.getId(), task.getImageUrl());
                    } else if (task.getStatus() == TaskStatus.FAILURE) {
                        failureTask.add(task);
                    }
                }

                // 失败的重跑
                if (!failureTask.isEmpty()) {
                    for (Task task : failureTask) {
                        SubmitImagineDTO imagineDTO = new SubmitImagineDTO();
                        imagineDTO.setPrompt(task.getPrompt());
                        Thread.sleep(IMAGINE_REQUEST_DELAY_MS); // Add delay before each imagine request
                        List<Prod> prods = prodMapper.selectByMjTaskId(task.getId());
                        // 更新所有相关记录
                        for (Prod prod : prods) {
                            prod.setAttritube1(null);
                            prod.setAttritube2(null);
                            prod.setAttritube3(null);
                            prod.setAttritube6(null);
                            prod.setPrompt(null);
                            prod.setTitle(null);
                        }

                        // 批量更新
                        prodMapper.batchUpdate(prods);

                    }
                }

                // 按taskId分组处理记录
                Map<String, List<Prod>> groupedProds = unprocessedProds.stream()
                        .collect(Collectors.groupingBy(Prod::getAttritube1));

                for (Map.Entry<String, List<Prod>> entry : groupedProds.entrySet()) {
                    String taskId = entry.getKey();
                    String imageUrl = taskIdToImageUrl.get(taskId);

                    if (imageUrl != null) {
                        List<Prod> prods = entry.getValue();
                        processImageForProds(imageUrl, prods, taskId);
                    }
                }
            }
        } catch (Exception e) {
            log.error("Error in file processing", e);
        } finally {
            isProcessing.set(false);
        }
    }


    private void processImageForProds(String imageUrl, List<Prod> prods, String taskId) {
        try {
            // 下载图片并处理WebP格式
            BufferedImage originalImage = null;
            try {
                originalImage = ImageIO.read(new URL(imageUrl));
                if (originalImage == null) {
                    log.error("Failed to read image from URL: {} - Image format may not be supported", imageUrl);
                    return;
                }
            } catch (IOException e) {
                log.error("Error downloading image from URL: {}", imageUrl, e);
                return;
            }

            int width = originalImage.getWidth();
            int height = originalImage.getHeight();

            // 计算每个部分的大小
            int partWidth = width / 2;
            int partHeight = height / 2;

            // 确保有足够的记录来存储所有部分
            if (prods.size() >= 4) {
                // 为当前任务创建目录
                String taskDir = BASE_PATH;
                new File(taskDir).mkdirs();

                // 分割图片并更新记录
                for (int i = 0; i < 4; i++) {
                    int row = i / 2;
                    int col = i % 2;

                    // 提取部分图片
                    BufferedImage part = originalImage.getSubimage(
                        col * partWidth,
                        row * partHeight,
                        partWidth,
                        partHeight
                    );

                    // 保存到本地文件系统
                    String fileName = String.format("%s-%d.png", prods.get(i).getDataId(), i + 1);
                    File outputFile = new File(SEL_PATH  + fileName);
                    ImageIO.write(part, "png", outputFile);

                    // 转换为字节数组并更新数据库
//                    byte[] imageData = convertImageToBytes(part);
                    Prod prod = prods.get(i);
//                    prod.setImageData(imageData);
                    // 保存本地文件路径到attritube3
                    prod.setAttritube3(taskDir + fileName);
                }

                // 批量更新数据库
                prodMapper.batchUpdateImagesWithPath(prods);
            }
        } catch (Exception e) {
            log.error("Error processing image from URL: " + imageUrl, e);
        }
    }

    private byte[] convertImageToBytes(BufferedImage image) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ImageIO.write(image, "png", baos);
        return baos.toByteArray();
    }
}